#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created by PyCharm.

@Date    : Thu Feb 25 2021 
@Time    : 09:17:41
@File    : eval.py
@Author  : alpha
"""


import cv2
import torch
import numpy as np
import matplotlib.pyplot as plt

from tqdm import tqdm
from pathlib import Path

from src.log import logger
from src.dataset import FaceSpoofEvalDataset
from src.model import FaceSpoofNet
from src.regvgg import RepVGG18Slim, RepVGG18Old


if __name__ == '__main__':

    # data_path = '/data/dataset/CelebA_Spoof_cropped/test/spoof/7806'
    data_path = '/data/dataset/CelebA_Spoof_cropped/motovis/test'
    # data_path = '/data/dataset/CelebA_Spoof_cropped/motovis_facespoof_selected'
    # model_path = 'checkpoints/repvgg18_slim/face_spoof_with_cls_best.pth'
    model_path = 'checkpoints/repvgg18_20210227_v2/face_spoof_with_cls_best.pth'
    save_path = Path('tmp')

    eval_dataset = FaceSpoofEvalDataset(root=data_path)
    model = FaceSpoofNet(feat_channels=256, backbone=RepVGG18Old)
    model.load_state_dict(torch.load(model_path))

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()

    with torch.no_grad():
        for im, im_path in tqdm(eval_dataset):
            im = im.to(device)[None]
            pred = model(im)[0]

            im = ((im[0] + 0.5).permute((1, 2, 0)).data.cpu().numpy() * 255).astype(np.uint8).copy()
            cv2.putText(
                im,
                ('spoof' if pred > 0.5 else 'live') + ': {:.2f}'.format(pred.item()),
                (16, 128),
                cv2.FONT_HERSHEY_DUPLEX,
                1,
                (250, 0, 0),
                thickness=2
            )

            save_img = save_path / Path(model_path).stem / im_path.relative_to(Path(data_path))
            save_img.parent.mkdir(parents=True, exist_ok=True)

            plt.imsave(save_img, im)

